import time
import cv2
from threading import Thread
from djitellopy import Tello

tello = Tello()

# 连接无人机
tello.connect()

keepRecording = True
tello.streamon()
frame_read = tello.get_frame_read()


def videoRecorder():
    # 获取正确的帧尺寸（注意宽度在前，高度在后）
    width = frame_read.frame.shape[1]  # 正确宽度
    height = frame_read.frame.shape[0]  # 正确高度

    # 使用更通用的 MP4 格式（需要安装视频编码器）
    video = cv2.VideoWriter(
        'tello_video.mp4',
        cv2.VideoWriter_fourcc(*'mp4v'),  # MP4 编码
        30,
        (width, height)  # 正确的尺寸顺序
    )

    while keepRecording:
        frame = frame_read.frame  # 获取当前帧
        video.write(frame)  # 写入视频
        time.sleep(1 / 30)  # 匹配 30fps 帧率

    video.release()


# 创建并启动视频录制线程
recorder = Thread(target=videoRecorder)
recorder.start()

# 飞行控制指令
try:
    time.sleep(5)
finally:
    # 确保无论是否出现异常都会停止录制
    keepRecording = False
    recorder.join()
    tello.streamoff()  # 关闭视频流
    cv2.destroyAllWindows()  # 关闭所有 OpenCV 窗口